#ifndef __CLASSIFACTION__
#define __CLASSIFACTION__

#include <opencv.hpp>
#include "DLKit.h"

extern "C"
{

	struct Classification
	{
		int class_id{ 0 };
		std::string className{ "A" };
		float confidence{ 0.0 };
		cv::Scalar color{};
		cv::Rect box{};
		cv::Mat Mask;
		float value[32] = { 100,160,160,6500 };
	};

	class VINORT_API ClassificationInfer
	{
	private:
		using HandleRT = void*;
		HandleRT core = nullptr, infer_request = nullptr, compiled_model = nullptr;
		double modelConfidenseThreshold = 0.25;
		double modelScoreThreshold = 0.25;
		double modelNMSThreshold = 0.45; //origin is 0.5
		std::vector<std::string> classes;
		int loadState = 1;
		cv::Size InferSize = cv::Size(224, 224);

	public:
		ClassificationInfer();
		~ClassificationInfer();

		std::vector<std::string> getClassesName();
		const char* getRuntimeType();
		std::string LoadModel(const char* fileName = "./DL/best.onnx", const char* labelName = "./DL/classes.txt", bool isGPU = true);
		std::vector<Classification> Infer(const cv::Mat& input, double r1 = -1, double c1 = -1, double r2 = -1, double c2 = -1);
		void setThreshold(double confidence = 0.25, double score = 0.25, double NMSthreshold = 0.5);
		void setInferSize(cv::Size sz);
		cv::Size getInferSize();
		void saveLabel(const cv::Mat & input, std::vector<Classification> detections, std::string fileName, std::string path = "./");

		enum {
			LOAD_OK = 0,
			LOAD_INFER_FILE_NOT_FOUND = 1,
			LOAD_TEXT_FILE_NOT_FOUND = 2
		};
		int LoadState();
	};
};
#endif
